from flask import Flask, request, send_file, make_response, jsonify
from your_image_generation_library import generate_image_from_parameters
from PIL import Image
import io  

from waitress import serve   


app = Flask(__name__)  

@app.route('/')  
def home():  
    return "Hello, World!"

@app.route('/api/images/generate', methods=['POST'])
def generate_image():  
    data = request.get_json()
    prompt = data.pop('prompt','best quality')  
    # image = data.get('image', None)  
    width = data.pop('width', None)  
    height = data.pop('height', None)  
    # num_inference_steps = data.get('num_inference_steps', 30)  
  
    # 假设generate_image_from_parameters返回一个文件路径或字节流  
    image = generate_image_from_parameters(prompt, width, height, **data)  
  
    # 将 PIL Image 对象转换为字节流  
    img_byte_arr = io.BytesIO()  
    image.save(img_byte_arr, format='JPEG')  
    img_byte_arr = img_byte_arr.getvalue()  
  
    # 设置响应类型为图像  
    response = make_response(img_byte_arr)  
    response.headers['Content-Type'] = 'image/jpeg'  
  
    return response


'''
# 调试用
if __name__ == '__main__':  
    app.run()
'''

# 部署服务
if __name__ == '__main__':  
    serve(app.wsgi_app, host='0.0.0.0', port=8000)

